package org.example.config;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.provider.AbstractDataSourceProvider;
import com.baomidou.dynamic.datasource.provider.DynamicDataSourceProvider;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.*;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.context.annotation.*;
import org.springframework.core.annotation.Order;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.Collections;
import java.util.List;
import java.util.Map;

@Configuration
@Import({DruidDynamicDataSourceConfiguration.class, DynamicDataSourceCreatorAutoConfiguration.class, DynamicDataSourceAopConfiguration.class, DynamicDataSourceAssistConfiguration.class})
@AutoConfigureBefore(
        value = {DataSourceAutoConfiguration.class},
        name = {"com.alibaba.druid.spring.boot.autoconfigure.DruidDataSourceAutoConfigure"}
)
@ConditionalOnProperty(
        prefix = "spring.datasource.dynamic",
        name = {"enabled"},
        havingValue = "true",
        matchIfMissing = true
)
public class DynamicShardingDataSourceConfig implements InitializingBean {

    @Value("${spring.datasource.dynamic.sharding.name}")
    String shardingName;

    private final DynamicDataSourceProperties properties;
    private final List<DynamicDataSourcePropertiesCustomizer> dataSourcePropertiesCustomizers;

    public DynamicShardingDataSourceConfig(DynamicDataSourceProperties properties, ObjectProvider<List<DynamicDataSourcePropertiesCustomizer>> dataSourcePropertiesCustomizers) {
        this.properties = properties;
        this.dataSourcePropertiesCustomizers = dataSourcePropertiesCustomizers.getIfAvailable();
    }

    @Lazy
    @Resource
    private DataSource shardingDataSource;

    @Bean("shardingDynamicDataSourceProvider")
    @Order(1)
    public DynamicDataSourceProvider shardingDynamicDataSourceProvider(DefaultDataSourceCreator defaultDataSourceCreator) {
        return new AbstractDataSourceProvider(defaultDataSourceCreator) {
            @Override
            public Map<String, DataSource> loadDataSources() {
                Map<String, DataSource> dataSourceMap = createDataSourceMap(Collections.emptyMap());
                dataSourceMap.put(shardingName, shardingDataSource);
                return dataSourceMap;
            }
        };
    }

    @Bean
    @Primary
    public DataSource dataSource(List<DynamicDataSourceProvider> providers) {
        DynamicRoutingDataSource dataSource = new DynamicRoutingDataSource(providers);
        dataSource.setPrimary(this.properties.getPrimary());
        dataSource.setStrict(this.properties.getStrict());
        dataSource.setStrategy(this.properties.getStrategy());
        dataSource.setP6spy(this.properties.getP6spy());
        dataSource.setSeata(this.properties.getSeata());
        dataSource.setGraceDestroy(this.properties.getGraceDestroy());
        return dataSource;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        if (!CollectionUtils.isEmpty(this.dataSourcePropertiesCustomizers)) {

            for (DynamicDataSourcePropertiesCustomizer customizer : this.dataSourcePropertiesCustomizers) {
                customizer.customize(this.properties);
            }
        }
    }
}
